#!/usr/bin/env python3
"""
标准人物LoRA训练参数节点（适配advanced_lora_trainer.py）
参照sd_scripts/docs/train_network_README-ja.md官方参数说明
"""

import folder_paths
import os
from .lora_trainer_utils.constants import CHARACTER_LORA_PARAMS_CATEGORY
class CharacterLoraParams:
    """标准人物LoRA训练参数节点"""

    aux_id = "comfyui_lora_train"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # 基础模型与数据
                "pretrained_model_name_or_path": (folder_paths.get_filename_list("checkpoints"), {
                    "default": "",
                    "tooltip": "预训练模型路径（如SD1.5/SDXL/写实底模）"
                }),
                "train_data_dir": ("STRING", {
                    "default": "",
                    "tooltip": "训练图片文件夹（建议子文件夹分人物/风格）"
                }),
                "output_dir": ("STRING", {
                    "default": "./output",
                    "tooltip": "模型输出目录"
                }),
                "output_name": ("STRING", {
                    "default": "character_lora",
                    "tooltip": "输出模型名称"
                }),
                # LoRA结构
                "network_dim": ("INT", {
                    "default": 32,
                    "min": 1,
                    "max": 128,
                    "tooltip": "LoRA网络维度（RANK），32-64推荐"
                }),
                "network_alpha": ("INT", {
                    "default": 16,
                    "min": 1,
                    "max": 128,
                    "tooltip": "LoRA alpha，建议为dim的1/2"
                }),
                "network_module": (["networks.lora", "networks.dylora", "networks.ia3", "networks.lokr", "networks.loha", "networks.c3liera", "networks.oft"], {
                    "default": "networks.lora",
                    "tooltip": "LoRA网络模块类型，默认使用networks.lora"
                }),
                # 训练核心参数
                "learning_rate": ("FLOAT", {
                    "default": 5e-5,
                    "min": 1e-6,
                    "max": 1e-3,
                    "tooltip": "学习率，人物推荐5e-5"
                }),
                "max_train_steps": ("INT", {
                    "default": 1000,
                    "min": 100,
                    "max": 20000,
                    "tooltip": "最大训练步数，人物训练推荐1000-2000，官方示例400"
                }),
                "train_batch_size": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 8,
                    "tooltip": "训练批次大小，显存允许可调高"
                }),
                "resolution": ("INT", {
                    "default": 512,
                    "min": 256,
                    "max": 1024,
                    "tooltip": "训练分辨率，人物建议512或768"
                }),
                # 优化器与调度器
                "optimizer_type": (["AdamW", "AdamW8bit", "PagedAdamW", "PagedAdamW8bit", "PagedAdamW32bit", "Lion8bit", "PagedLion8bit", "Lion", "SGDNesterov", "SGDNesterov8bit", "DAdaptation", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptAdanIP", "DAdaptLion", "DAdaptSGD", "AdaFactor"], {
                    "default": "AdamW",
                    "tooltip": "优化器类型：AdamW(默认), AdamW8bit(需要bitsandbytes), Lion, SGDNesterov等"
                }),
                "lr_scheduler": (["cosine_with_restarts", "constant_with_warmup", "linear", "cosine"], {
                    "default": "cosine_with_restarts",
                    "tooltip": "学习率调度器"
                }),
                # 训练特性
                "mixed_precision": (["fp16", "bf16", "no"], {
                    "default": "fp16",
                    "tooltip": "混合精度训练"
                }),
                "cache_latents": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "缓存Latents，加速训练"
                }),
                "gradient_checkpointing": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "梯度检查点，节省显存"
                }),
                "clip_skip": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 12,
                    "tooltip": "CLIP跳过层数，人物训练建议2"
                }),
                # 保存与中断
                "save_every_n_steps": ("INT", {
                    "default": 500,
                    "min": 100,
                    "max": 5000,
                    "tooltip": "每多少步保存一次模型"
                }),
                "save_model_as": (["safetensors", "ckpt"], {
                    "default": "safetensors",
                    "tooltip": "模型保存格式"
                }),
            },
            "optional": {
                "config_file": ("STRING", {
                    "default": "",
                    "tooltip": "可选TOML配置文件"
                }),
                "vae_name": (folder_paths.get_filename_list("vae"), {
                    "default": "",
                    "tooltip": "可选VAE模型路径"
                }),
                # 训练模块选择
                "network_train_unet_only": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "仅训练UNet，不训练文本编码器"
                }),
                "network_train_text_encoder_only": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "仅训练文本编码器，不训练UNet"
                }),
                # 差别化学习率
                "unet_lr": ("FLOAT", {
                    "default": 0.0,
                    "min": 0.0,
                    "max": 1e-3,
                    "tooltip": "UNet学习率，0表示使用主学习率"
                }),
                "text_encoder_lr": ("FLOAT", {
                    "default": 0.0,
                    "min": 0.0,
                    "max": 1e-3,
                    "tooltip": "文本编码器学习率，0表示使用主学习率，建议5e-5"
                }),
                # 数据加载优化
                "max_data_loader_n_workers": ("INT", {
                    "default": 8,
                    "min": 0,
                    "max": 16,
                    "tooltip": "数据加载器工作进程数，显存不足时调低"
                }),
                "persistent_data_loader_workers": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "持久化数据加载器工作进程，Windows环境推荐"
                }),
                # 高级训练选项
                "alpha_mask": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "使用图像alpha通道作为遮罩，透明图像训练"
                }),
                "network_weights": ("STRING", {
                    "default": "",
                    "tooltip": "预训练LoRA权重路径，继续训练用"
                }),
                # 学习率预热
                "lr_warmup_steps": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 1000,
                    "tooltip": "学习率预热步数"
                }),
                # 梯度累积
                "gradient_accumulation_steps": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 16,
                    "tooltip": "梯度累积步数，显存不足时调高"
                }),
                # 随机种子
                "seed": ("INT", {
                    "default": 42,
                    "min": 0,
                    "max": 2147483647,
                    "tooltip": "随机种子"
                }),
            }
        }
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        """验证输入参数 - 只验证路径是否存在"""
        # 检查训练数据目录是否存在
        train_data_dir = kwargs.get("train_data_dir")
        if train_data_dir and isinstance(train_data_dir, str) and train_data_dir.strip() and not os.path.exists(train_data_dir):
            return f"训练数据目录不存在: {train_data_dir}"
        
        # 检查采样提示词文件（如果提供）
        sample_prompts = kwargs.get("sample_prompts")
        if sample_prompts and isinstance(sample_prompts, str) and sample_prompts.strip() and not os.path.exists(sample_prompts):
            return f"采样提示词文件不存在: {sample_prompts}"
        
        return True
    RETURN_TYPES = ("TRAINING_PARAMS",)
    RETURN_NAMES = ("训练参数",)
    FUNCTION = "generate_params"
    CATEGORY = CHARACTER_LORA_PARAMS_CATEGORY

    def generate_params(self, **kwargs):
        # 自动处理所有常见路径参数，转为绝对路径
        path_keys = [
            ("pretrained_model_name_or_path", "checkpoints"),
            ("vae_name", "vae"),
            ("config_file", None),
            ("network_weights", None),
        ]
        
        for key, folder in path_keys:
            if kwargs.get(key):
                try:
                    if folder:
                        kwargs[key] = folder_paths.get_full_path(folder, kwargs[key])
                    else:
                        if not os.path.isabs(kwargs[key]):
                            kwargs[key] = os.path.abspath(kwargs[key])
                except Exception as e:
                    print(f"[WARNING] 获取{key}完整路径失败: {e}")
        # 处理分辨率格式
        if "resolution" in kwargs:
            resolution = kwargs["resolution"]
            kwargs["resolution"] = f"{resolution},{resolution}"
        # 输出目录
        if kwargs.get("output_dir"):
            output_dir = kwargs["output_dir"]
            if not os.path.isabs(output_dir):
                output_dir = os.path.abspath(output_dir)
            kwargs["output_dir"] = output_dir
            try:
                os.makedirs(output_dir, exist_ok=True)
                print(f"[INFO] 输出目录已创建/确认: {output_dir}")
            except Exception as e:
                print(f"[WARNING] 创建输出目录失败: {e}")
        # 训练数据目录
        if kwargs.get("train_data_dir"):
            train_data_dir = kwargs["train_data_dir"]
            if not os.path.isabs(train_data_dir):
                train_data_dir = os.path.abspath(train_data_dir)
            kwargs["train_data_dir"] = train_data_dir
            if not os.path.exists(train_data_dir):
                print(f"[WARNING] 训练数据目录不存在: {train_data_dir}")
        # 正则化数据目录
        if kwargs.get("reg_data_dir"):
            reg_data_dir = kwargs["reg_data_dir"]
            if not os.path.isabs(reg_data_dir):
                reg_data_dir = os.path.abspath(reg_data_dir)
            kwargs["reg_data_dir"] = reg_data_dir
            if not os.path.exists(reg_data_dir):
                print(f"[WARNING] 正则化数据目录不存在: {reg_data_dir}")
        # 差别化学习率
        if kwargs.get("unet_lr") == 0.0:
            kwargs.pop("unet_lr", None)
        if kwargs.get("text_encoder_lr") == 0.0:
            kwargs.pop("text_encoder_lr", None)
        return (kwargs,)

# 节点注册
NODE_CLASS_MAPPINGS = {
    "CharacterLoraParams": CharacterLoraParams
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "CharacterLoraParams": "人物LoRA训练参数"
}